{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetune and deploy Cohere Command Light Model from AWS Bedrock\n",
    "\n",
    "This sample notebook shows you how to finetune and deploy cohere command light models using Amazon SageMaker.\n",
    "\n",
    "> **Note**: This is a reference notebook and it cannot run unless you make changes suggested in the notebook.\n",
    "\n",
    "## Pre-requisites:\n",
    "1. **Note**: This notebook contains elements which render correctly in Jupyter interface. Open this notebook from an Amazon SageMaker Notebook Instance or Amazon SageMaker Studio.\n",
    "1. Ensure that you have an IAM role set up for bedrock model customization https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-iam-role.html\n",
    "\n",
    "## Contents:\n",
    "1. [Subscribe to Amazon Bedrock](#1.-Subscribe-to-Amazon-Bedrock)\n",
    "2. [Run the model customization job](#2.-Run-the-model-customization-job)\n",
    "   1. [Upload training data](#A.-Upload-training-data)\n",
    "   2. [Finetune models on uploaded data](#B.-Finetune-models-on-uploaded-data)\n",
    "   3. [Wait for job to complete](#C.-Wait-for-job-to-complete)\n",
    "3. [Create an endpoint for inference](#3.-Create-an-endpoint-for-inference)\n",
    "   1. [Provision model throughput](#A.-Provision-model-throughput)\n",
    "   2. [Perform real-time inference](#B.-Perform-real-time-inference)\n",
    "4. [Clean-up](#4.-Clean-up)\n",
    "    1. [Delete the endpoint](#A.-Delete-the-endpoint)    \n",
    "\n",
    "## Usage instructions\n",
    "You can run this notebook one cell at a time (By using Shift+Enter for running a cell)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Subscribe to Amazon Bedrock"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Follow the instructions in the [Amazon Bedrock](https://console.aws.amazon.com/bedrock) console and [add access](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) to the cohere command light model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Run the model customization job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade cohere-aws\n",
    "# if you upgrade the package, you need to restart the kernel\n",
    "\n",
    "import cohere_aws\n",
    "import boto3"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Upload training data\n",
    "\n",
    "Choose a directory on S3 to store the training data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_bucket_name = \"finetune-data\" # bucket where data should be uploaded to, your bedrock model customization IAM role should also have access to this bucket\n",
    "s3_train_data_path = \"generative/train.jsonl\" # the path where train data will be stored\n",
    "s3_eval_data_path = \"generative/eval.jsonl\" # the path where eval data will be stored (optional)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload sample training data and optional eval data to S3:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3 = boto3.client('s3')\n",
    "\n",
    "# upload example data to s3\n",
    "s3.upload_file(\"../examples/sample_generative_data.jsonl\", s3_bucket_name, s3_train_data_path)\n",
    "s3.upload_file(\"../examples/sample_generative_data_eval.jsonl\", s3_bucket_name, s3_eval_data_path) # (optional)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Finetune models on uploaded data\n",
    "\n",
    "Specify a directory on S3 where finetuned models should be stored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_models_dir = \"s3://finetuned_models/\" # where the models will be saved"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Cohere client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "co = cohere_aws.Client(mode=cohere_aws.Mode.BEDROCK, region_name=region)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the fine-tuning job:\n",
    "> **Note**: Update the role ARN with the role created in step 2 of the pre-requisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_url = f\"s3://{s3_bucket_name}/{s3_train_data_path}\"\n",
    "eval_data_url = f\"s3://{s3_bucket_name}/{s3_eval_data_path}\"\n",
    "\n",
    "job_id = co.create_finetune(\n",
    "    name=\"finetuned-model\",\n",
    "    base_model_id=\"cohere.command-light-text-v14:7:4k\", \n",
    "    train_data=train_data_url,\n",
    "    s3_models_dir=s3_models_dir,\n",
    "    eval_data=eval_data_url, \n",
    "    role=\"arn:aws:iam::<ACCOUNT_ID>:role/service-role/<ROLE_NAME>\"\n",
    ")\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Wait for job to complete\n",
    "\n",
    "> **Note**: This job may take a while to complete, if it does not complete within the timeout your job will still be running and you may need to wait longer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = co.wait_for_finetune_job(job_id)\n",
    "print(model_id)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create an endpoint for inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Provision model throughput\n",
    "\n",
    "The Cohere AWS SDK provides a built-in method for provisioning throughput to inference finetuned models.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_arn = co.provision_throughput(model_id=model_id, name=\"custom-model-throughput\", model_units=1)\n",
    "print(model_arn)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Perform real-time inference\n",
    "\n",
    "Now, you can access all models deployed on the endpoint for inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = co.generate(prompt=\"hello\", model_id=model_arn)\n",
    "print(result)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Clean-up"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Delete the endpoint\n",
    "\n",
    "After you've successfully performed inference, you can delete the deployed endpoint to avoid being charged continuously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock = boto3.client(\"bedrock\", region_name=region)\n",
    "bedrock.delete_provisioned_model_throughput(model_arn)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "emb",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
